
import nni
from torch.optim import Adam
import utils
# from tensorboardX import SummaryWriter
import time
# from utils import cprint
from dataloader import check_dataset, getDataset
from models import check_model, getModel
from loss import getLoss
from Procedure import TRAIN, TEST, TEST_with_auc
from pprint import pprint
import os
import torch

def main(config):

    print('===========config================')
    pprint(config)
    print('===========end===================')
    utils.set_seed(config["seed"])
    print(">>SEED:", config["seed"])

    check_dataset(config["dataset"])
    check_model(config["model"])

    dataset = getDataset(config["data_path"], config["dataset"])
    model = getModel(config["model"], config, dataset).to(config["device"])
    loss_fun = getLoss(config["loss"], model, config, dataset)
    opt = Adam(model.parameters(), lr=config['lr'])
    try:
        early_stopping = utils.EarlyStopping(patience=5, verbose=config["is_use_early_stop"], path=utils.getFileName(config))
        for epoch in range(config["TRAIN_epochs"]):
            results = TEST(dataset, model, epoch, None, config['multicore'])
            early_stopping(results["ndcg"][2], model)
            if config["is_use_early_stop"] and early_stopping.early_stop:
                print("Early stopping")
                break
            nni.report_intermediate_result(results["ndcg"][2])
            output_information = TRAIN(dataset, model, loss_fun, opt, epoch, w=None)
            print(f'EPOCH[{epoch + 1}/{config["TRAIN_epochs"]}] {output_information}')
        nni.report_final_result(early_stopping.best_score)
        # final test
        model.load_state_dict(torch.load(utils.getFileName(config)))
        results = TEST_with_auc(dataset, model, None, config['multicore'])
    finally:
        print("code over!!!")

if __name__ == "__main__":
    import Config
    main(Config.config)
